from app import Strategy, Context, ResultAnalyse
from pydantic import BaseModel
import time

class GridParameter(BaseModel):
    instrument: str # 交易标的代码
    marketId: int # 市场类型
    basePrice: float # 基准价
    offsetPriceBuy: float # 买入价差
    offsetRatioBuy: float # 买入价差百分比
    offsetPriceSell: float # 卖出价差
    offsetRatioSell: float # 卖出价差百分比
    quantityBuy: int # 每次买入数量
    balanceBuy: float # 每次买入金额
    quantitySell: int # 每次卖出数量
    balanceSell: float # 每次卖出金额
    priceHighLimit: float # 价格区间——最高价
    priceLowLimit: float # 价格区间——最低价
    bidPricingLevel: int # 使用市场买盘档位 - 默认买一
    askPricingLevel: int # 使用市场卖盘档位 - 默认卖一
    expireTime: int # 策略有效期, 时间戳
    increaseEntrust: bool # 是否翻倍开仓
    maxFund: float # 初始最大可用资金 - 动态调整
    maxQtyCanSell: int # 初始最大可卖量 - 动态调整
    fillWhenJump: bool # 倍数委托
    pullbackRatio: float # 回落卖出幅度
    reBounceRatio: float # 反弹买入幅度
    t0: bool # 是否是T0标的
    
class GridInfo(BaseModel):
    priceBuy: float = 0 # 买入价格
    priceSell: float = 0 # 卖出价格
    pauseBuy: bool = False # 买入暂停
    pauseSell: bool = False # 卖出暂停
    
class GridStrategy(Strategy.Strategy):
    
    def __init__(self, name, params, context: Context):
        super().__init__(name, params, context)
        result = self.__check_params__()
        if result:
            self.status = 1 # 策略状态，1-开始，0-结束，-1-异常
            self.order_id = 0
            print("start success")
        else:
            self.status = -1
            print("start failed")

    def __reset_grid_price__(self, price):
        self.grid_info = GridInfo()
        self.grid_params.basePrice = price
        if self.grid_params.offsetPriceBuy > 0:
            self.grid_info.priceBuy = self.grid_params.basePrice - self.grid_params.offsetPriceBuy
        else:
            self.grid_info.priceBuy = self.grid_params.basePrice * (1-self.grid_params.offsetRatioBuy)

        if self.grid_params.offsetPriceSell > 0:
            self.grid_info.priceSell = self.grid_params.basePrice + self.grid_params.offsetPriceSell
        else:
            self.grid_info.priceSell = self.grid_params.basePrice * (1+self.grid_params.offsetRatioSell)
        self.grid_info.pauseBuy = False
        self.grid_info.pauseSell = False
        return True
    
    def __check_params__(self):
        self.grid_params = GridParameter(**self.params)
        # 检查参数是否有效
        print(self.grid_params)
        result = self.__reset_grid_price__(self.grid_params.basePrice)
        return result

    def __calculate_buy_qty__(self, last_price):
        qty = self.grid_params.quantityBuy
        if qty == 0:
            qty = (int)(self.grid_params.balanceBuy / last_price)
        return qty
    
    def __calculate_sell_qty__(self, last_price):
        qty = self.grid_params.quantitySell
        if qty == 0:
            qty = (int)(self.grid_params.balanceSell / last_price)
        return qty
    
    def on_daily_close(self):
        current_asset = 0
        for key,value in self.context.account.position.items():
            current_asset += value.qty * value.last_price
        current_asset += self.context.account.money
        self.context.account.assets[self.context.account.date] = current_asset
        print(f"{self.context.account.date}, {current_asset}")
        print(self.context.account.position)
    
    def on_snapshot(self, data: Context.Snapshot):
        if self.grid_params.instrument != data.code:
            print("wrong code")
            return
        if self.context.has_position(data.code):
            current_position = self.context.account.position[data.code]
            current_position.last_price = data.last_price
            self.context.account.position[data.code] = current_position
        if self.context.account.date != data.date:
            self.context.account.date = data.date
        if self.status != 1:
            print("has reach final status")
            return
        print(data)
        if data.last_price > self.grid_params.priceHighLimit or data.last_price < self.grid_params.priceLowLimit:
            print("reach grid limit price, ignore")
            return
        if self.grid_info.priceBuy > data.last_price and not self.grid_info.pauseBuy: # 买入
            self.order_id += 1
            qty = self.__calculate_buy_qty__(data.last_price)
            if qty > 0:
                order = Context.Order(order_id=self.order_id, code=data.code, market=data.market,direction=1, price=data.last_price, qty=qty, date=data.date, time=data.time)
                result = self.context.order(order)
                if result:
                    self.__reset_grid_price__(data.last_price)
                else:
                    self.grid_info.pauseBuy = True
        elif self.grid_info.priceSell < data.last_price and not self.grid_info.pauseSell: # 卖出
            self.order_id += 1
            qty = self.__calculate_sell_qty__(data.last_price)
            if qty > 0:
                order = Context.Order(order_id=self.order_id, code=data.code, market=data.market,direction=0, price=data.last_price, qty=qty, date=data.date, time=data.time)
                result = self.context.order(order)
                if result:
                    self.__reset_grid_price__(data.last_price)
                else:
                    self.grid_info.pauseSell = True
        
    def on_complete(self):
        analyse = ResultAnalyse.ResultAnalyse()
        analyse.draw_result_html(self.context.account)
        self.status = 0